package com.example.aigc_education.ws;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.example.aigc_education.common.BusinessException;
import com.example.aigc_education.domain.ChatHistory;
import com.example.aigc_education.domain.request.ChatAiTeacherRequest;
import com.example.aigc_education.enums.ChatHistoryType;
import com.example.aigc_education.mapper.ChatHistoryMapper;
import com.example.aigc_education.security.domain.LoginUser;
import com.example.aigc_education.utils.JsonUtils;
import com.example.aigc_education.utils.StringUtils;
import com.example.aigc_education.utils.WebSocketUtils;
import io.github.lnyocly.ai4j.listener.SseListener;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletion;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatCompletionResponse;
import io.github.lnyocly.ai4j.platform.openai.chat.entity.ChatMessage;
import io.github.lnyocly.ai4j.service.IChatService;
import io.github.lnyocly.ai4j.service.PlatformType;
import io.github.lnyocly.ai4j.service.factor.AiService;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import javax.annotation.Resource;
import java.io.Serializable;
import java.security.Principal;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
public class AiTeacherWSHandler extends AbstractWebSocketHandler {

    @Resource
    private ChatHistoryMapper chatHistoryMapper;

    @Resource
    private AiService aiService;

    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        WebSocketSessionHolder.addSession(session.getId(), session);
    }

    /**
     * 处理发送来的文本消息
     * @param session    会话
     * @param message    消息
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String sessionId = session.getId();
        String payload = message.getPayload();
        log.info("[receive message] sessionId: {} , message:{}", sessionId, payload);


        try{
            ChatAiTeacherRequest chatAiTeacherRequest = JsonUtils.fromJson(payload, ChatAiTeacherRequest.class);
            if (chatAiTeacherRequest == null) {
                log.warn("无效的请求格式, sessionId: {}", session.getId());
                WebSocketUtils.sendMessage(session, "无效的请求格式");
                return;
            }

            LambdaQueryWrapper<ChatHistory> queryWrapper = new LambdaQueryWrapper<>();
            queryWrapper.eq(ChatHistory::getBotId, chatAiTeacherRequest.getBotId());
            List<ChatMessage> history = chatHistoryMapper.selectList(queryWrapper)
                    .stream()
                    .map(hs -> {

                        if ("system".equals(hs.getRole())){
                            return ChatMessage.withSystem(hs.getContent());
                        }

                        if ("user".equals(hs.getRole())){
                            return ChatMessage.withUser(hs.getContent());
                        }

                        return ChatMessage.withAssistant(hs.getContent());
                    })
                    .collect(Collectors.toList());

            if (history.size() == 1){
                history.add(ChatMessage.withUser("你好"));
                IChatService chatService = aiService.getChatService(PlatformType.MINIMAX);
                ChatCompletion chatCompletion = ChatCompletion.builder()
                        .model("minimax-text-01")
                        .messages(history)
                        .build();

                ChatCompletionResponse response = chatService.chatCompletion(chatCompletion);
                ChatMessage chatMessage = response.getChoices().get(0).getMessage();
                ChatHistory aiHistory = new ChatHistory();
                aiHistory.setBotId(chatAiTeacherRequest.getBotId());
                aiHistory.setContent(chatMessage.getContent().getText());
                aiHistory.setRole(ChatHistoryType.ASSISTANT.getTitle());
                chatHistoryMapper.insert(aiHistory);
                history.add(chatMessage);
            }

            history.add(ChatMessage.withUser(chatAiTeacherRequest.getContent()));
            IChatService chatService = aiService.getChatService(PlatformType.MINIMAX);
            ChatCompletion chatCompletion = ChatCompletion.builder()
                    .model("minimax-text-01")
                    .messages(history)
                    .build();

            // 构造监听器
            final AiMessage[] aiMessage = {null};
            SseListener sseListener = new SseListener() {
                @Override
                protected void send() {
                    aiMessage[0] = new AiMessage();
                    aiMessage[0].setContent(this.getCurrStr());
                    aiMessage[0].setRole(ChatHistoryType.ASSISTANT.getTitle());
                    aiMessage[0].setConvId(chatAiTeacherRequest.getBotId().toString());
                    aiMessage[0].setMessageId(String.valueOf(System.currentTimeMillis()));
                    WebSocketUtils.sendMessage(session, JsonUtils.toJson(aiMessage[0]));
                }
            };
            // 显示函数参数，默认不显示
            sseListener.setShowToolArgs(true);

            // 发送SSE请求
            chatService.chatCompletionStream(chatCompletion, sseListener);


            // 保存用户输入
            ChatHistory chatHistory = new ChatHistory();
            chatHistory.setBotId(chatAiTeacherRequest.getBotId());
            chatHistory.setContent(chatAiTeacherRequest.getContent());
            chatHistory.setRole(ChatHistoryType.USER.getTitle());
            chatHistoryMapper.insert(chatHistory);

            // 保存AI回复
            ChatHistory aiHistory = new ChatHistory();
            aiHistory.setBotId(chatAiTeacherRequest.getBotId());
            aiHistory.setContent(String.valueOf(sseListener.getOutput()));
            aiHistory.setRole(ChatHistoryType.ASSISTANT.getTitle());
            chatHistoryMapper.insert(aiHistory);

            aiMessage[0].setFinished(true);
            WebSocketUtils.sendMessage(session, JsonUtils.toJson(aiMessage[0]));
            log.info("[send message] sessionId: {}, message: {}", sessionId, sseListener.getOutput());


        }catch (Exception e){
            log.error("[send message error] sessionId: {}, message: {}", sessionId, e.getMessage());
        }


    }


    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        log.info("[receive binary message] sessionId: {} , message:{}", session.getId(), message.getPayload());
        super.handleBinaryMessage(session, message);
    }

    /**
     * 心跳监测的回复
     *
     * @param session   会话
     * @param message   消息
     * @throws Exception
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        WebSocketUtils.sendPongMessage(session);
    }
    /**
     * 连接出错时
     *
     * @param session   会话
     * @param exception 异常
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[transport error] sessionId: {} , exception:{}", session.getId(), exception.getMessage());
        if (session.isOpen()) {
            session.close(CloseStatus.SERVER_ERROR);
        }
    }

    /**
     * 连接关闭后
     *
     * @param session   会话
     * @param status    关闭状态
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        WebSocketSessionHolder.removeSession(session.getId());
    }

    /**
     * 指示处理程序是否支持接收部分消息
     *
     * @return 如果支持接收部分消息，则返回true；否则返回false
     */
    @Override
    public boolean supportsPartialMessages() {
        return super.supportsPartialMessages();
    }

    @Data
    public static class AiMessage implements Serializable {
        private String content;
        private String role;
        private String messageId;
        private String convId;
        private Boolean finished = false;
    }
}
